import obspy as ob
from scipy import signal
from pylab import *
from matplotlib import rcParams
rcParams['font.family'] = 'Times New Roman'
rcParams['font.size'] = '8'
from STATs import *
from multiprocessing import Pool
from scipy import interpolate
from obspy.signal.filter import bandpass

# 默认参数
ABS_AMP=False
DO_PAC = True
NORM = True
UP_RATE=1 # 不做升采样
DT_DEFAULT=0.02
EPS=1e-10

# 归一函数
def norm(d,EPS=EPS):
    d = d-np.mean(d)
    d = d/(np.max(d)+EPS)
    return d

# 时窗滑动平均函数
def trace_do_pac(trace, N=25, EPS=EPS):

    nt = len(trace)
    data = trace[:int(nt//N*N)].reshape([N,nt//N])
    d_amp = np.sqrt(np.mean(data**2, axis=1))*2
    step = (np.arange(0,N)+0.5) * int(nt//N)
    t = np.arange(0,nt)
    f = interpolate.interp1d(step, d_amp,kind='slinear',bounds_error=False, fill_value=(d_amp[0],d_amp[-1]))
    mask = f(t)
    if len(np.where(mask<0)[0])>0:
        raise ValueError('in interpolate, amp value < 0',len(np.where(mask<0)[0]),step,d_amp)
    trace = trace/(mask+EPS)
    return trace
 
# # 叠加单一台站, 旧版本，复杂，功能多
def stacked_traces_1stat_old(data_1stat,FS, dt = DT_DEFAULT, UP_RATE=UP_RATE, DO_PAC=DO_PAC,DO_SELECT = False,DO_WHITEN=False, **paras):
    data_correlate = data_1stat.copy()
    ne,nt = data_correlate.shape
    DO_FILTER = True
    fs, fe = FS
    weights = np.ones([ne]) if paras.get('event_weights') is None else paras.get('event_weights')
    NORM = False if paras.get('NORM') is None else paras.get('NORM')
    assert weights.size==ne

    f = 1/dt/2
    b,a = signal.butter(4,[fs/f,fe/f],'bandpass')
    if DO_PAC and DO_SELECT:
        print('ONLY PAC')
    if DO_PAC or DO_SELECT:
        for ie in range(ne):
            trace = data_correlate[ie,:]
            if DO_FILTER:
                trace = signal.filtfilt(b, a, trace)
            # 时移平滑
            if DO_PAC:
                trace =trace_do_pac(trace, N=int(nt*dt)*2)
            else:
                if DO_SELECT:
                    E = np.mean(trace**2)**0.5 # 平均振幅
                    if E<0.1:
                        trace = 0
            # 谱白化
            if DO_WHITEN:
                df = 1/(dt*len(trace))
                smooth_nf = int(0.4/df)//2*2+1 # 0.2Hz平滑窗, odd
                spec = np.fft.fft(trace)
                spec_amp = np.abs(spec)
                spec_amp = signal.savgol_filter(spec_amp,smooth_nf,1)
                spec_amp = spec_amp/(1e-2+np.max(spec_amp))
                spec = spec/(1e-2+spec_amp)
                trace = np.real(np.fft.ifft(spec))
                trace = signal.filtfilt(b, a, trace)

            data_correlate[ie,:] = trace*weights[ie]
        return data_correlate.sum(axis=0)
    else:
        result = 0
        for ie in range(ne):
            trace = data_correlate[ie,:]
            if DO_FILTER:
                trace = signal.filtfilt(b, a, trace)
            if NORM:
                trace = norm(trace)
            result = result+trace
        return result

# # 
def stacked_traces_1stat(    data_1stat,FS, dt = DT_DEFAULT, UP_RATE=UP_RATE, DO_PAC=DO_PAC, NORM=NORM):        
    '''
    叠加单一台站所有记录
    data_1stat: [ne,nt]
    '''
    
    ne,nt = data_1stat.shape
    t = np.arange(nt)*dt
    dt_new = dt/UP_RATE
    t_new = np.arange(nt*UP_RATE)*dt_new
    
    fs, fe = FS
    result = 0

    f_max = 1/dt/2
    b,a = signal.butter(4,[fs/f_max,fe/f_max],'bandpass')

    for ie in range(ne):
        trace = data_1stat[ie,:]
        trace =signal.filtfilt(b,a,trace)
        # trace = bandpass(trace, freqmin=fs, freqmax=fe, df=1/dt, corners=4, zerophase=True)
        # 叠加前升采样
        if UP_RATE!=1:
            f = interpolate.interp1d(t,trace,bounds_error=False, fill_value=0)
            trace =f(t_new)
        # 时窗平滑
        if DO_PAC:
            trace = trace_do_pac(trace)
        if NORM:
            trace = norm(trace)
        result = result+trace

    # print(data_1stat.shape)
    f_max_new = 1/dt_new/2
    b,a = signal.butter(4,[fs/f_max_new,fe/f_max_new],'bandpass')
    result = signal.filtfilt(b,a,result)
    # result = bandpass(result, freqmin=fs, freqmax=fe, df=1/dt_new, corners=4, zerophase=True) # 太慢了

    return result[::UP_RATE]  


def stacked_traces_ns(data_all_stats, N_Core=12, FS = [1,12],dt=DT_DEFAULT, UP_RATE=UP_RATE, DO_PAC=DO_PAC, NORM=NORM):
    '''
    子函数， 对三维数据[ne,ns,nt] ne 维度并行叠加
    data:[ne,ns,nt]
    N_Core: 核心数目
    FS: 带通滤波
    dt: dt
    UP_RATE: 叠前上采样倍率，默认为1
    DO_PAC: 是否时窗滑动平均
    NORM: 叠后归一化
    '''
    # print(FS, dt,UP_RATE, DO_PAC)
    ne,ns,nt = data_all_stats.shape
    input_args = []
    for i in range(ns):
        data_1stat_i = data_all_stats[:,i, :]
        input_args.append([data_1stat_i,FS, dt,UP_RATE, DO_PAC])
        
    with Pool(N_Core) as pool:
        results = pool.starmap(stacked_traces_1stat,input_args)

    traces_PAC = np.zeros([ns,nt])

    for i in range(ns):
        if not NORM:
            traces_PAC[i,:] = results[i]
        else:
            traces_PAC[i,:] = norm(results[i])
    return traces_PAC

def stacked_traces_ns_nopar(data_all_stats,FS = [1,12],dt=DT_DEFAULT, UP_RATE=UP_RATE, DO_PAC=DO_PAC, NORM=NORM, **paras):
    '''
    子函数， 对三维数据[ne,ns,nt] ne 维度叠加, 非并行
    data:[ne,ns,nt]
    FS: 带通滤波
    dt: dt
    UP_RATE: 叠前上采样倍率，默认为1
    DO_PAC: 是否时窗滑动平均
    NORM: 叠后归一化
    '''
    ne,ns,nt = data_all_stats.shape
    traces = np.zeros([ns,nt])
    for i in range(ns):
        data_1stat_i = data_all_stats[:,i, :]
        trace_stacked = stacked_traces_1stat(data_1stat_i,FS, dt,UP_RATE, DO_PAC)

        if not NORM:
            traces[i,:] = trace_stacked
        else:
            traces[i,:] = norm(trace_stacked)
    return traces